// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

#nullable enable

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;

namespace DanmakuR.Protocol.Buffer.Writers;

internal sealed class MemoryBufferWriter : Stream, IBufferWriter<byte>
{
	[ThreadStatic]
	private static MemoryBufferWriter? _cachedInstance;

#if DEBUG
	private bool _inUse;
#endif

	private readonly int _minimumSegmentSize;
	private int _bytesWritten;

	private List<CompletedBuffer>? _completedSegments;
	private byte[]? _currentSegment;
	private int _position;

	public MemoryBufferWriter(int minimumSegmentSize = 4096)
	{
		_minimumSegmentSize = minimumSegmentSize;
	}

	public override long Length => _bytesWritten;
	public override bool CanRead => false;
	public override bool CanSeek => false;
	public override bool CanWrite => true;
	public override long Position
	{
		get => throw new NotSupportedException();
		set => throw new NotSupportedException();
	}

	public static MemoryBufferWriter Get()
	{
		var writer = _cachedInstance;
		if (writer == null)
		{
			writer = new MemoryBufferWriter();
		}
		else
		{
			// Taken off the thread static
			_cachedInstance = null;
		}
#if DEBUG
		if (writer._inUse)
		{
			throw new InvalidOperationException("The reader wasn't returned!");
		}

		writer._inUse = true;
#endif

		return writer;
	}

	public static void Return(MemoryBufferWriter writer)
	{
		_cachedInstance = writer;
#if DEBUG
		writer._inUse = false;
#endif
		writer.Reset();
	}

	public void Reset()
	{
		if (_completedSegments != null)
		{
			for (var i = 0; i < _completedSegments.Count; i++)
			{
				_completedSegments[i].Return();
			}

			_completedSegments.Clear();
		}

		if (_currentSegment != null)
		{
			ArrayPool<byte>.Shared.Return(_currentSegment);
			_currentSegment = null;
		}

		_bytesWritten = 0;
		_position = 0;
	}

	public void Advance(int count)
	{
		_bytesWritten += count;
		_position += count;
	}

	public Memory<byte> GetMemory(int sizeHint = 0)
	{
		EnsureCapacity(sizeHint);

		return _currentSegment.AsMemory(_position, _currentSegment.Length - _position);
	}

	public Span<byte> GetSpan(int sizeHint = 0)
	{
		EnsureCapacity(sizeHint);

		return _currentSegment.AsSpan(_position, _currentSegment.Length - _position);
	}

	public void CopyTo(IBufferWriter<byte> destination)
	{
		if (_completedSegments != null)
		{
			// Copy completed segments
			var count = _completedSegments.Count;
			for (var i = 0; i < count; i++)
			{
				destination.Write(_completedSegments[i].Span);
			}
		}

		destination.Write(_currentSegment.AsSpan(0, _position));
	}

	public override Task CopyToAsync(Stream destination, int bufferSize, CancellationToken cancellationToken)
	{
		if (_completedSegments == null && _currentSegment is not null)
		{
			// There is only one segment so write without awaiting.
			return destination.WriteAsync(_currentSegment, 0, _position, cancellationToken);
		}

		return CopyToSlowAsync(destination, cancellationToken);
	}

	[MemberNotNull(nameof(_currentSegment))]
	private void EnsureCapacity(int sizeHint)
	{
		// This does the Right Thing. It only subtracts _position from the current segment length if it's non-null.
		// If _currentSegment is null, it returns 0.
		var remainingSize = _currentSegment?.Length - _position ?? 0;

		// If the sizeHint is 0, any capacity will do
		// Otherwise, the buffer must have enough space for the entire size hint, or we need to add a segment.
		if (sizeHint == 0 && remainingSize > 0 || sizeHint > 0 && remainingSize >= sizeHint)
		{
			// We have capacity in the current segment
#pragma warning disable CS8774 // Member must have a non-null value when exiting.
			return;
#pragma warning restore CS8774 // Member must have a non-null value when exiting.
		}

		AddSegment(sizeHint);
	}

	[MemberNotNull(nameof(_currentSegment))]
	private void AddSegment(int sizeHint = 0)
	{
		if (_currentSegment != null)
		{
			// We're adding a segment to the list
			if (_completedSegments == null)
			{
				_completedSegments = new List<CompletedBuffer>();
			}

			// Position might be less than the segment length if there wasn't enough space to satisfy the sizeHint when
			// GetMemory was called. In that case we'll take the current segment and call it "completed", but need to
			// ignore any empty space in it.
			_completedSegments.Add(new CompletedBuffer(_currentSegment, _position));
		}

		// Get a new buffer using the minimum segment size, unless the size hint is larger than a single segment.
		_currentSegment = ArrayPool<byte>.Shared.Rent(Math.Max(_minimumSegmentSize, sizeHint));
		_position = 0;
	}

	private async Task CopyToSlowAsync(Stream destination, CancellationToken cancellationToken)
	{
		if (_completedSegments != null)
		{
			// Copy full segments
			var count = _completedSegments.Count;
			for (var i = 0; i < count; i++)
			{
				var segment = _completedSegments[i];
#if NETCOREAPP
				await destination.WriteAsync(segment.Buffer.AsMemory(0, segment.Length), cancellationToken).ConfigureAwait(false);
#else
				await destination.WriteAsync(segment.Buffer, 0, segment.Length, cancellationToken).ConfigureAwait(false);
#endif
			}
		}

		if (_currentSegment is not null)
		{
#if NETCOREAPP
			await destination.WriteAsync(_currentSegment.AsMemory(0, _position), cancellationToken).ConfigureAwait(false);
#else
			await destination.WriteAsync(_currentSegment, 0, _position, cancellationToken).ConfigureAwait(false);
#endif
		}
	}

	public byte[] ToArray()
	{
		if (_currentSegment == null)
		{
			return Array.Empty<byte>();
		}

		var result = new byte[_bytesWritten];

		var totalWritten = 0;

		if (_completedSegments != null)
		{
			// Copy full segments
			var count = _completedSegments.Count;
			for (var i = 0; i < count; i++)
			{
				var segment = _completedSegments[i];
				segment.Span.CopyTo(result.AsSpan(totalWritten));
				totalWritten += segment.Span.Length;
			}
		}

		// Copy current incomplete segment
		_currentSegment.AsSpan(0, _position).CopyTo(result.AsSpan(totalWritten));

		return result;
	}

	public void CopyTo(Span<byte> span)
	{
		Debug.Assert(span.Length >= _bytesWritten);

		if (_currentSegment == null)
		{
			return;
		}

		var totalWritten = 0;

		if (_completedSegments != null)
		{
			// Copy full segments
			var count = _completedSegments.Count;
			for (var i = 0; i < count; i++)
			{
				var segment = _completedSegments[i];
				segment.Span.CopyTo(span.Slice(totalWritten));
				totalWritten += segment.Span.Length;
			}
		}

		// Copy current incomplete segment
		_currentSegment.AsSpan(0, _position).CopyTo(span.Slice(totalWritten));

		Debug.Assert(_bytesWritten == totalWritten + _position);
	}

	public override void Flush() { }
	public override Task FlushAsync(CancellationToken cancellationToken) => Task.CompletedTask;
	public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
	public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
	public override void SetLength(long value) => throw new NotSupportedException();

	public override void WriteByte(byte value)
	{
		if (_currentSegment != null && (uint)_position < (uint)_currentSegment.Length)
		{
			_currentSegment[_position] = value;
		}
		else
		{
			AddSegment();
			_currentSegment[0] = value;
		}

		_position++;
		_bytesWritten++;
	}

	public override void Write(byte[] buffer, int offset, int count)
	{
		var position = _position;
		if (_currentSegment != null && position < _currentSegment.Length - count)
		{
			System.Buffer.BlockCopy(buffer, offset, _currentSegment, position, count);

			_position = position + count;
			_bytesWritten += count;
		}
		else
		{
			BuffersExtensions.Write(this, buffer.AsSpan(offset, count));
		}
	}

#if NETCOREAPP
	public override void Write(ReadOnlySpan<byte> span)
	{
		if (_currentSegment != null && span.TryCopyTo(_currentSegment.AsSpan(_position)))
		{
			_position += span.Length;
			_bytesWritten += span.Length;
		}
		else
		{
			BuffersExtensions.Write(this, span);
		}
	}
#endif

	public WrittenBuffers DetachAndReset()
	{
		_completedSegments ??= new List<CompletedBuffer>();

		if (_currentSegment is not null)
		{
			_completedSegments.Add(new CompletedBuffer(_currentSegment, _position));
		}

		var written = new WrittenBuffers(_completedSegments, _bytesWritten);

		_currentSegment = null;
		_completedSegments = null;
		_bytesWritten = 0;
		_position = 0;

		return written;
	}

	public WrittenSequence DeatchToSequence()
	{
		_completedSegments ??= new List<CompletedBuffer>();

		if (_currentSegment is not null)
		{
			_completedSegments.Add(new CompletedBuffer(_currentSegment, _position));
		}

		var seq = new WrittenSequence(_completedSegments, _bytesWritten);

		_currentSegment = null;
		_completedSegments = null;
		_bytesWritten = 0;
		_position = 0;

		return seq;
	}

	protected override void Dispose(bool disposing)
	{
		if (disposing)
		{
			Reset();
		}
	}

	/// <summary>
	/// Holds the written segments from a MemoryBufferWriter and is no longer attached to a MemoryBufferWriter.
	/// You are now responsible for calling Dispose on this type to return the memory to the pool.
	/// </summary>
	internal readonly ref struct WrittenBuffers
	{
		public readonly List<CompletedBuffer> Segments;
		private readonly int _bytesWritten;

		public WrittenBuffers(List<CompletedBuffer> segments, int bytesWritten)
		{
			Segments = segments;
			_bytesWritten = bytesWritten;
		}

		public int ByteLength => _bytesWritten;

		public void Dispose()
		{
			for (var i = 0; i < Segments.Count; i++)
			{
				Segments[i].Return();
			}
			Segments.Clear();
		}
	}

	/// <summary>
	/// Holds the written segments from a MemoryBufferWriter and is no longer attached to a MemoryBufferWriter.
	/// You are now responsible for calling Dispose on this type to return the memory to the pool.
	/// </summary>
	internal readonly struct WrittenSequence : IDisposable
	{
		private sealed class Segment : ReadOnlySequenceSegment<byte>
		{
			private readonly byte[] _buffer;
			private readonly int _end;
			private Segment? _next;

			public int End => _end;

			public Segment(byte[] buffer, int end)
			{
				_buffer = buffer;
				_end = end;

				Memory = buffer.AsMemory(0, end);
			}

			public Segment? NextSegment
			{
				get => _next;
				set
				{
					_next = value;
					Next = value;
				}
			}

			public Segment CreateNext(byte[] buffer, int end)
			{
				NextSegment = new(buffer, end);
				NextSegment.RunningIndex += _end;

				return NextSegment;
			}
		}

		private readonly byte[]? singleSegment;
		private readonly List<CompletedBuffer>? Segments;
		public int ByteLength { get; }
		internal WrittenSequence(List<CompletedBuffer> segments, int bytesWritten)
		{
			singleSegment = null;
			Segments = segments;
			ByteLength = bytesWritten;
		}

		internal WrittenSequence(byte[] buffer, int bytesWritten)
		{
			Segments = null;
			singleSegment = buffer;
			ByteLength = bytesWritten;
		}

		public void Dispose()
		{
			if (Segments == null)
			{
				Debug.Assert(singleSegment != null);
				ArrayPool<byte>.Shared.Return(singleSegment);
			}	
			else
			{
				for (int i = 0; i < Segments.Count; i++)
				{
					Segments[i].Return();
				}
				Segments.Clear(); 
			}
		}

		public ReadOnlySequence<byte> GetSequence()
		{
			if (singleSegment != null)
				return new(singleSegment, 0, ByteLength);

			Debug.Assert(Segments != null);
			CompletedBuffer current;
			if (Segments.Count == 1)
			{
				current = Segments[0];
				return new(current.Buffer, 0, current.Length);
			}
			else if (Segments.Count == 0)
			{
				return default;
			}
			else
			{
				Segment first = new(Segments[0].Buffer, Segments[0].Length);
				Segment? segment = null;
				for (int i = 1; i < Segments.Count; i++)
				{
					current = Segments[i];
					segment = segment?.CreateNext(current.Buffer, current.Length)
						?? first.CreateNext(current.Buffer, current.Length);
				}

				return new(first, first.End, segment!, segment!.End);
			}
		}
	}

	/// <summary>
	/// Holds a byte[] from the pool and a size value. Basically a Memory but guaranteed to be backed by an ArrayPool byte[], so that we know we can return it.
	/// </summary>
	internal readonly struct CompletedBuffer
	{
		public byte[] Buffer { get; }
		public int Length { get; }

		public ReadOnlySpan<byte> Span => Buffer.AsSpan(0, Length);

		public CompletedBuffer(byte[] buffer, int length)
		{
			Buffer = buffer;
			Length = length;
		}

		public void Return()
		{
			ArrayPool<byte>.Shared.Return(Buffer);
		}
	}
}
